#-*- coding=utf-8 -*-
import partie1 as p1
import matplotlib.pyplot as mp
import numpy as np

#------------------------------------------------------------------------#
#                              PENDULE SIMPLE                            #
#------------------------------------------------------------------------#

def plot_sol_pendulum(t0,y0,f):
    "Affiche la solution approximée sur l' intervalle [0 10]"
    mp.clf()
    nmax = 300
    x = np.arange(0.,10.+10./nmax,10./nmax)
    res = p1.meth_n_step(y0, t0, nmax, 10./nmax, f,p1.step_kutta)[0]
    y1 = []
    y2 = []
    for i in range(len(res)):
        y1 = y1 + [res[i][0][0]]
        y2 = y2 + [res[i][1][0]]
    g = mp.plot(x, y1, linewidth=1.0)
    h = mp.plot(x, y2, linewidth=1.0)
    mp.legend((g,h),("Angle theta","Vitesse angulaire"))
    mp.show()

def calc_freq(theta):
    "Calcule la fréquence pour un angle initial theta donné"
    t0 = np.array([[0.]])
    y0 = np.array([[theta],[0.]])
    g = 10
    l = 1
    f = lambda X,T: np.array([[X[1][0]],[(-g/l)*np.sin(X[0][0])]])
    nmax = 500
    x = np.arange(0.,10.+10./nmax,10./nmax)
    a = abs(theta)/theta
    res = p1.meth_n_step(y0, t0, nmax, 10./nmax, f,p1.step_kutta)[0]
    i = 10
    while((abs(res[i][1][0])>abs(theta)/10. or a*res[i][0][0]<0) and i < len(res)-1):
        i = i+1
    return 1./x[i]

def plot_freq():
    "Trace la fréquence en fonction de l'angle initial theta"
    theta = np.arange(-2*np.pi/3,2*np.pi/3,0.2)
    y = []
    for i in range(len(theta)):
        y = y + [calc_freq(theta[i])]
    g = mp.plot(theta, y, linewidth=1.0)
    mp.xlabel("Angle initial")
    mp.ylabel("Frequence du pendule")
    mp.show()

def pendulum1(theta):
    "Résoud l'équation différentielle en fonction de l'angle initial theta"
    t0 = np.array([[0.]])
    y0 = np.array([[theta],[0.]])
    g = 10
    l = 1
    f = lambda X,T: np.array([[X[1][0]],[(-g/l)*np.sin(X[0][0])]])
    plot_sol_pendulum(t0,y0,f)

#------------------------------------------------------------------------#
#                              PENDULE DOUBLE                            #
#------------------------------------------------------------------------#

def plot_sol_pendulum2(t0,y0,f):
    "Affiche la solution approximée sur l' intervalle [0 10]"
    mp.clf()
    nmax = 500
    x = np.arange(0.,10.+10./nmax,10./nmax)
    res = p1.meth_n_step(y0, t0, nmax, 10./nmax, f,p1.step_kutta)[0]
    y1 = []
    y2 = []
    y3 = []
    y4 = []
    for i in range(len(res)):
        y1 = y1 + [res[i][0][0]]
        y2 = y2 + [res[i][1][0]]
        y3 = y3 + [res[i][2][0]]
        y4 = y4 + [res[i][3][0]]
    g = mp.plot(x, y1, linewidth=1.0)
    h = mp.plot(x, y2, linewidth=1.0)
    i = mp.plot(x, y3, linewidth=1.0)
    j = mp.plot(x, y4, 'y',linewidth=1.0)
    mp.legend((g,h,i,j),("Angle theta1","Angle theta2","Vitesse angulaire 1","Vitesse angulaire 2"))
    mp.show()

def plot_ext_pendulum2(t0,y0,f,l1,l2):
    "Affiche la positiion de l'extrémité du pendule au cours du temps"
    mp.clf()
    nmax = 1000
    res = p1.meth_n_step(y0, t0, nmax, 10./nmax, f,p1.step_kutta)[0]
    res2 = p1.meth_n_step(y0+np.array([[0.1],[0.],[0.],[0.]]), t0, nmax, 10./nmax, f,p1.step_kutta)[0]
    y1 = []#theta1
    y2 = []#theta2
    x = []
    y = []
    for i in range(len(res)):
        y1 = y1 + [res[i][0][0]]
        y2 = y2 + [res[i][1][0]]
    for i in range(len(res)):
        x = x + [l1*np.sin(y1[i])-l2*np.sin(y2[i])]
        y = y + [-l1*np.cos(y1[i])-l2*np.cos(y2[i])]
    g = mp.plot(x, y, linewidth=1.0)
    y1 = []#theta1
    y2 = []#theta2
    x = []
    y = []
    for i in range(len(res)):
        y1 = y1 + [res2[i][0][0]]
        y2 = y2 + [res2[i][1][0]]
    for i in range(len(res)):
        x = x + [l1*np.sin(y1[i])-l2*np.sin(y2[i])]
        y = y + [-l1*np.cos(y1[i])-l2*np.cos(y2[i])]
    h = mp.plot(x, y, linewidth=1.0)
    mp.legend((g,h),("theta1 = Pi/5, theta2 = 2Pi/3","theta1 = Pi/5+0.1, theta2 = 2Pi/3"))
    mp.show()


def pendulum2(theta1,theta2):
    "Résoud l'équation différentielle en fonction des angles initiaux theta1 et theta2"
    t0 = np.array([[0.]])
    y0 = np.array([[theta1],[theta2],[0.],[0.]])
    g = 10
    l1 = 1
    l2 = 1
    m1 = 1
    m2 = 1
    f = lambda X,T: np.array([[X[2][0]],[X[3][0]],[(-g*(2*m1+m2)*np.sin(X[0][0])-m2*g*np.sin(X[0][0]-2*X[1][0])-2*np.sin(X[0][0]-X[1][0])*m2*(l2*X[3][0]**2+l1*np.cos(X[0][0]-X[1][0])*X[2][0]**2))/(l1*(2*m1+m2-m2*np.cos(2*X[0][0]-2*X[1][0])))],[(2*np.sin(X[0][0]-X[1][0])*((m1+m2)*l1*X[2][0]**2+g*(m1+m2)*np.cos(X[0][0])+l2*m2*np.cos(X[0][0]-X[1][0])*X[3][0]**2))/(l1*(2*m1+m2-m2*np.cos(2*X[0][0]-2*X[1][0])))]])
    plot_sol_pendulum2(t0,y0,f)
    plot_ext_pendulum2(t0,y0,f,l1,l2)

def meth_n_step(y0, t0, n, h, f, meth = p1.step_euler):
    "Calcule la suite de n point pour une méthode donnée"
    res = [y0]
    t = t0
    for i in range(n):
        res.append(meth(res[i], t, h, f))
        t = t + h
    return res


def time_flip(x,res):
    "Calcule le temps du premier retournement pour une courbe donnée"
    i = 1
    while(np.pi-abs(res[i][0][0])>0. and np.pi-abs(res[i][1][0])>0. and i < len(res)-1):
        i = i+1
    return i
    
def plot_time(n):
    g = 10
    l1 = 1
    l2 = 1
    m1 = 1
    m2 = 1
    f = lambda X,T: np.array([[X[2][0]],[X[3][0]],[(-g*(2*m1+m2)*np.sin(X[0][0])-m2*g*np.sin(X[0][0]-2*X[1][0])-2*np.sin(X[0][0]-X[1][0])*m2*(l2*X[3][0]**2+l1*np.cos(X[0][0]-X[1][0])*X[2][0]**2))/(l1*(2*m1+m2-m2*np.cos(2*X[0][0]-2*X[1][0])))],[(2*np.sin(X[0][0]-X[1][0])*((m1+m2)*l1*X[2][0]**2+g*(m1+m2)*np.cos(X[0][0])+l2*m2*np.cos(X[0][0]-X[1][0])*X[3][0]**2))/(l1*(2*m1+m2-m2*np.cos(2*X[0][0]-2*X[1][0])))]])
    nmax = 100
    tmax = 10.
    x = np.arange(0.,tmax+tmax/nmax,tmax/nmax)
    A = np.zeros([n,n])
    a = (2*np.pi)/A.shape[0]
    b = -np.pi
    for i in range(A.shape[0]):
        for j in range(A.shape[1]):
            print ((i*A.shape[1]+j)/(float(A.shape[0]*A.shape[1])))*100.,"%"
            t0 = np.array([[0.]])
            y0 = np.array([[i*a+b],[j*a+b],[0.],[0.]])
            res =  meth_n_step(y0, t0, nmax, 10./nmax, f,p1.step_kutta)
            k = time_flip(x,res)
            if(k == len(x)-1):
                A[i][j] = x[0]
            else:
                A[i][j] = x[time_flip(x,res)]
    mp.clf()
    mp.imshow(A)
    mp.gcf()
    mp.clim()
    mp.colorbar()
    mp.show()


pendulum1(-np.pi/2.)
plot_freq()
pendulum2(np.pi/5.,2*np.pi/3.)
plot_time(30)
